package org.wyz.spark;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.broadcast.Broadcast;
import scala.Tuple2;

import java.util.Arrays;
import java.util.Map;

public class KMeans {
    private static final String[] ARGS = {"data/input/kmeans/point.txt", "data/input/kmeans/center.txt", "data/output/kmeans"};
    private static final int EPOCH = 5;

    public static void main(String[] args) {
        SparkConf conf = new SparkConf().setMaster("local").setAppName("wyz-kmeans");
        JavaSparkContext sc = new JavaSparkContext(conf);


        JavaRDD<Point> points = sc.textFile(ARGS[0]).map(a -> {
            String[] ss = a.split("\\s")[0].split(",");
            return new Point(Double.parseDouble(ss[0]), Double.parseDouble(ss[1]));
        });


        Map<String, Point> rawCenters = sc.textFile(ARGS[1]).mapToPair(a -> {
            String[] ss = a.split("\\s");
            String[] sp = ss[1].split(",");
            System.out.println(Arrays.asList(sp));
            return new Tuple2<>(ss[0], new Point(Double.parseDouble(sp[0]), Double.parseDouble(sp[1])));
        }).collectAsMap();

        System.out.println(rawCenters);
//        sc.stop();
//        System.exit(1);

        Broadcast<Map<String, Point>> broadcast = sc.broadcast(rawCenters);

        for (int i = 1; i <= EPOCH; i++) {
            Broadcast<Map<String, Point>> finalBroadcast = broadcast;
            JavaPairRDD<String, Point> centerRDD = points.mapToPair(a -> {
                        Map<String, Point> centers = finalBroadcast.value();
                        String minTag = null;
                        double minDis = Double.MAX_VALUE;
                        for (String tag : centers.keySet()) {
                            double dis = a.distance(centers.get(tag));
                            if (dis < minDis) {
                                minTag = tag;
                                minDis = dis;
                            }
                        }
                        return new Tuple2<>(minTag, new Tuple2<>(a, 1));
                    }).reduceByKey((a, b) -> new Tuple2<>(a._1.sum(b._1), a._2 + b._2))
                    .mapToPair(a -> new Tuple2<>(a._1, a._2._1.mean(a._2._2)));

            if (i == EPOCH) {
                centerRDD.map(a -> a._1 + " " + a._2).saveAsTextFile(ARGS[2]);
            } else {
                // 更新广播变量
                broadcast = sc.broadcast(centerRDD.collectAsMap());
            }
        }
        sc.stop();
    }
}
